import networkx as nx
import matplotlib.pyplot as plt
from collections import deque
from pulp import *
from tools import generate_gas_network, visualize_network, get_outward_connected_components
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径
        
def visualize(H, pos):
    G = H.copy()
    for u in ["super_source", "super_sink"]:
        if u in G.nodes():
            G.remove_node(u)
    plt.figure(figsize=(14, 10))
    # 按类型定义节点颜色
    node_colors = []
    node_sizes = []
    edgecolors = []
    linewidths = []
    nodes = list(G.nodes)  # 创建静态的节点列表快照
    for node in nodes:
        try:
            attr = H.nodes[node]
            if attr['type'] == 'source':
                node_colors.append('lime')
                node_sizes.append(300)
            elif 'leak' in attr['type']:
                node_colors.append('red')
                node_sizes.append(300)                   
            elif 'user' in attr['type']:
                node_colors.append('gold')
                node_sizes.append(300)         
            elif attr['type'] == 'other':
                node_colors.append('green')
                node_sizes.append(300)       
            else:
            #elif attr['type'] == 'valve':
                node_colors.append('blue')
                node_sizes.append(300)
        except Exception as exp:
            node_colors.append('blue')
            node_sizes.append(300)            
    # 绘制节点
    nx.draw_networkx_nodes(H, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9)
    plt.show()
    exit()

    node_colors = []
    node_sizes = []
    edgecolors = []
    linewidths = []
    H = G.copy()
    nodes = list(H.nodes)  # 创建静态的节点列表快照
    for node in nodes:
        attr = H.nodes[node]
        if attr['type'] == 'valve' and node in cut_nodes:
            node_colors.append('blue')
            node_sizes.append(300)
            edgecolors.append("gold")
            linewidths.append(2)
        else:
            H.remove_node(node)
    # 绘制节点
    nx.draw_networkx_nodes(H, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9, edgecolors = 'red', linewidths=3)
  
    # 绘制边
    edge_colors = []
    edge_widths = []
    for u, v in G.edges():
        if (u, v) not in no_reachable_edges:
            edge_colors.append('gray')
            edge_widths.append(1.5)
            
        else:
            edge_colors.append('red')
            edge_widths.append(1.5)
    nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, arrows=True, arrowstyle='-|>', arrowsize=20)

    if len(no_reachable_edges) > 0:
        ax = plt.gca()  # 获取当前坐标轴
        for u, v in no_reachable_edges:
            x1, y1 = pos[u]
            x2, y2 = pos[v]
            mid_x = (x1 + x2) / 2
            mid_y = (y1 + y2) / 2
            ax.plot(mid_x, mid_y, 'rx', markersize=15, markeredgewidth=2)  # 红色叉号
    if len(leak_impact_edges) > 0:
        print ("leak_impact_edges", leak_impact_edges)
        ax = plt.gca()  # 获取当前坐标轴            
        for u, v in leak_impact_edges:
            x1, y1 = pos[u]
            x2, y2 = pos[v]
            mid_x = (x1 + x2) / 2
            mid_y = (y1 + y2) / 2
            ax.plot(mid_x, mid_y, 'rx', markersize=15, markeredgewidth=2)  # 红色叉号

    # if len(min_cut_set) > 0:
    #     edge_colors = []
    #     edge_widths = []
    #     edge_styles = []
    #     for u, v in G.edges():
    #         if (u, v) in min_cut_set:
    #             edge_colors.append('red')
    #             edge_widths.append(1.5)
    #             edge_styles.append('dashed')
        
    #     nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, style=edge_styles, arrows=True, arrowstyle='-|>')
  
    # 节点标签
    labels = {}
    for node in G.nodes:
        attr = G.nodes[node]
        if attr['type'] == 'valve':
            labels[node] = f"V{node}"
        elif attr['type'] == 'user':
            labels[node] = f"U{node}"
        elif attr['type'] == 'leak':
            labels[node] = f"L{node}"
        elif attr['type'] == 'other':
            labels[node] = f"O{node}"            
        elif attr['type'] == 'source':
            labels[node] = f"S{node}"
        else:
            labels[node] = f"{node}"
  
    nx.draw_networkx_labels(G, pos, labels, font_size=9)
  
    # 管道属性标注
    for (u, v), attr in G.edges.items():
        plt.text((pos[u][0]+pos[v][0])/2, 
                 (pos[u][1]+pos[v][1])/2, 
                 f"{G[u][v].get('capacity')}", 
                 fontsize=7, color='blue')
  
    # 图例
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', label='气源(S)', markerfacecolor='lime', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='用户(U)', markerfacecolor='gold', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='泄漏点(L)', markerfacecolor='red', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='其他节点(O)', markerfacecolor='green', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(V)未关', markerfacecolor='blue', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(V)关闭', markerfacecolor='blue', markeredgecolor='red', markeredgewidth=3, markersize=10),
        #plt.Line2D([0], [0], marker='o', color='w', label='增加节点(超级源/超级汇)', markerfacecolor='blue', markersize=10),
        plt.Line2D([0], [0], color='gray', lw=1, label='普通管道'),
        plt.Line2D([0], [0], marker='x', color='red', linestyle='None', markersize=15, markeredgewidth=2, label='被阻断的管道')

    ]
  
    plt.legend(handles=legend_elements, loc='upper right', prop=font)
    plt.title(name, fontsize=14, fontproperties=font)
    plt.tight_layout()
    plt.savefig('%s.jpg' %name, dpi=300)
    plt.show()

def find_min_cut(G, super_source, super_sink, valves):
    """基于流向的有向图最小割算法"""
    # 计算最小割
    min_cut_valve, partition = nx.minimum_cut(G, super_source, super_sink)
    #print ("!!!!!!!!!!!")
    #print ("minimum_cut", nx.minimum_cut(G, super_source, super_sink))
    #print ("maximum_flow", nx.maximum_flow(G, super_source, super_sink))
    print ("partition", partition)
    print ("min_cut_valve", min_cut_valve)
    reachable, non_reachable = partition

    min_cut_set = []
    for u, v in G.edges():
        if u in reachable and v in non_reachable:
            if G.has_edge(u, v) and u in valves:
                min_cut_set.append((u, v))
    return min_cut_set

if __name__ == "__main__":
    show = False
    #生成图
    G_ori, pos = generate_gas_network()

    if show:
        nx.draw_networkx(G_ori, pos)
        plt.show()

    #predecessors
    #successors
    #generate_cost_flow_network(G)

    #节点
    leaks = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'leak']
    sources = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'source']
    valves = [n for n, attr in G_ori.nodes(data=True) if attr['type'] == 'valve']

    #加入超级源和超级汇
    G = G_ori.copy()
    super_source = "super_source"
    super_sink = "super_sink"
    pos[super_source] = [0,0.3]
    pos[super_sink] = [0,-0.3]
    nodes = []
    nodes.append((super_source, {'type': 'super_source'}))
    nodes.append((super_sink, {'type': 'super_sink'}))
    G.add_nodes_from(nodes)
    for u in sources:
        G.add_edge(super_source, u)
    for u in leaks:
        G.add_edge(u, super_sink)
    
    if show:
        nx.draw_networkx(G, pos)
        plt.show()
    for u, v in G.edges():
        G[u][v]['capacity'] = 1e9
        G[u][v]['weight'] = 0

    #拆点（user节点和leak节点）
    H = G.copy()
    for u in G.nodes():
        if G.nodes[u]['type'] == 'user':
            u_out = f"{u}_out"
            H.add_nodes_from([u_out])
            for v in G.successors(u):
                H.remove_edge(u, v)
                H.add_edge(u_out, v, capacity = 1e9, weight = 0)
            pos[u_out] = pos[u] + [0.02, 0.02]
            H.add_edge(u, u_out, capacity = 1e9, weight = 1)
        elif G.nodes[u]['type'] == 'leak':
            u_out = f"{u}_out"
            H.add_nodes_from([u_out])
            for v in G.successors(u):
                H.remove_edge(u, v)
                H.add_edge(u_out, v, capacity = 1e9, weight = 0)
            pos[u_out] = pos[u] + [0.02, 0.02]
            H.add_edge(u, u_out, capacity = 1, weight = 0)
    
    if show:
        nx.draw_networkx(H, pos)
        plt.show()
    
    #最小费用最大流（保证用户数最小、泄漏数最多的流量）
    result = nx.max_flow_min_cost(H, super_source, super_sink)
    print (result)

    #加入超级源和超级汇，容量=上述流量，求最小割
    G = G_ori.copy()
    nodes = []
    nodes.append((super_source, {'type': 'super_source'}))
    nodes.append((super_sink, {'type': 'super_sink'}))
    G.add_nodes_from(nodes)
    for u in valves:
        G.add_edge(super_source, u)
    for u in leaks:
        G.add_edge(u, super_sink)
    for u, v in G.edges():
        if u in result and v in result[u]:# and u in valves:
            G[u][v]['capacity'] = result[u][v]
        else:
            G[u][v]['capacity'] = 1e9
    result2 = find_min_cut(G, super_source, super_sink, valves)
    print (result2)
    
